#!/usr/bin/python

import unittest

from hips.scheduling import problem
from hips.scheduling import schedule_mocks


class ProblemTestCase(unittest.TestCase):

  def testCreateVariables(self):
    mock_schedule = schedule_mocks.CreateMockSchedule(days=3)
    p = problem.SchedulingProblem(mock_schedule)
    v = p.CreateVariables(3, [0, 2], [1, 3])
    self.assertEquals(4, len(v))
    self.assertEquals([3, 5, 6], v[0])
    self.assertEquals([1, 2, 4], v[1])
    self.assertEquals([3, 5, 6], v[2])
    self.assertEquals([1, 2, 4], v[3])


if __name__ == '__main__':
  unittest.main()
